import time

from langgraph.constants import END
from langgraph.graph import StateGraph
from src.nl2sql.state import Nl2SqlState
from src.nl2sql.nodes import (start_entry,
                              retrieve_main_table,
                              choice_configuration_tables, generate_sql, should_find_tables, continue_find
                              )




nl2sql_graph_builder = StateGraph(Nl2SqlState)
nl2sql_graph_builder.add_node("start_entry", start_entry)
# 后续更新表注解之后，retrieve_main_table节点需要修改
nl2sql_graph_builder.add_node("retrieve_main_table", retrieve_main_table)
nl2sql_graph_builder.add_node("choice_configuration_tables", choice_configuration_tables)
nl2sql_graph_builder.add_node("continue_find", continue_find)
nl2sql_graph_builder.add_node("generate_sql", generate_sql)

nl2sql_graph_builder.set_entry_point("start_entry")
nl2sql_graph_builder.add_edge("start_entry", "retrieve_main_table")
nl2sql_graph_builder.add_edge("retrieve_main_table", "choice_configuration_tables")
nl2sql_graph_builder.add_conditional_edges("choice_configuration_tables", should_find_tables, ["generate_sql", "continue_find"])
nl2sql_graph_builder.add_edge("continue_find", "generate_sql")
nl2sql_graph = nl2sql_graph_builder.compile()


if __name__ == '__main__':
    start = time.time()
    response = nl2sql_graph.stream(
        {"nl2sql_messages": [{"role": "system", "content": "你是人类的好帮手"},
                             ("human", "查询上周出勤管理情况")]},  # 获取运行科2025年1月份的污泥外运情况#查询上周出勤管理情况#查询药剂入库情况
        {"recursion_limit": 10},

    )

    # print(response["main_table"])
    # print(response["configuration_tables"])
    # print(response["sql"])
    # print(time.time()-start)
    # for message in response["nl2sql_messages"]:
    #     string_representation = f"{message.type.upper()}: {message.content}\n"
    #     print(string_representation)
